package com.z1.core.kvutils.rocksdb;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.log4j.Logger;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
import org.rocksdb.Options;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksIterator;

import com.z1.core.PathUtil;
import com.z1.core.io.FileCreator;
import com.z1.core.io.ObjectByte;
import com.z1.core.kvutils.KvFactoryConst;
import com.z1.core.kvutils.KvObject;

public final class KvRocksDb implements KvObject {
	
	private final Logger logger = Logger.getLogger(getClass());
	private final String KV_PATH = new File(PathUtil.getConfDir(), "rockdb").getAbsolutePath();
	
	private Map<String, ColumnFamilyHandle> groupMap = null;
	
	private RocksDB rocksDb;
	
	static {
		try {
			RocksDB.loadLibrary();//加载jni
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
	}

	private KvRocksDb() {
		logger.info("KvRocksDb path = " + KV_PATH);
		//打开数据库
		try {
			List<ColumnFamilyDescriptor> columnFamilyDescriptors = initGroups();
	        List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
	        DBOptions dbOptions = new DBOptions();
	        dbOptions.setCreateIfMissing(true);
	        FileCreator.dirCreator(PathUtil.getConfDir());
			rocksDb = RocksDB.open(dbOptions, KV_PATH, columnFamilyDescriptors, columnFamilyHandles);
			initColumnFamilyHandle(columnFamilyHandles);
		} catch (Throwable e) {
			throw new RuntimeException(e);
		}
		Runtime.getRuntime().addShutdownHook(new Thread() {
			@Override
			public void run() {
				if (null != rocksDb) {
					try { rocksDb.close(); }
					catch (Throwable e) { }
				}
			}
		});
	}
	
	private List<ColumnFamilyDescriptor> initGroups() {
		List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
		try {
			Options rockOpts = new Options();
			rockOpts.setCreateIfMissing(true);
			//获取组
			groupMap = new HashMap<String, ColumnFamilyHandle>();
			List<byte[]> cfs = RocksDB.listColumnFamilies(rockOpts, KV_PATH);
			if (cfs.size() > 0) {
				for (byte[] cf : cfs) {
					ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(cf, new ColumnFamilyOptions());
					columnFamilyDescriptors.add(columnFamilyDescriptor);
				}
			} else {
				ColumnFamilyDescriptor columnFamilyDescriptor = new ColumnFamilyDescriptor(GROUP_DEFAULT.getBytes(), new ColumnFamilyOptions());
				columnFamilyDescriptors.add(columnFamilyDescriptor);
			}
		}
		catch (Throwable e) {
			throw new RuntimeException(e);
		}
		return columnFamilyDescriptors;
	}
	
	private void initColumnFamilyHandle(List<ColumnFamilyHandle> columnFamilyHandles) {
		try {
			for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
				groupMap.put(
					new String(columnFamilyHandle.getName()), 
					columnFamilyHandle
				);
			}
		}
		catch (Throwable e) {
			throw new RuntimeException(e);
		}
	}
	
	@Override
	public void set(String group, String key, Object value) {
		if (null == rocksDb) {
			return;
		} 
		if (null == group) {
			group = GROUP_DEFAULT;
		}
		try {
			ColumnFamilyHandle columnFamilyHandle = groupMap.get(group);
			if (null == columnFamilyHandle) {
				columnFamilyHandle = rocksDb.createColumnFamily(
					new ColumnFamilyDescriptor(
						group.getBytes(), 
						new ColumnFamilyOptions()
					)
				);
				groupMap.put(group, columnFamilyHandle);
			}
			rocksDb.put(
				columnFamilyHandle, 
				key.getBytes(), 
				null == value ? null : ObjectByte.toBytes(value)
			); 
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
	}
	
	@Override
	public List<String> groups() {
		List<String> groups = new ArrayList<String>();
		for (String key : groupMap.keySet()) {
			groups.add(key);
		}
		return groups;
	}

	@Override
	public Map<String, Object> get(String group) {
		if (null == group) {
			group = GROUP_DEFAULT;
		}
		Map<String, Object> map = new HashMap<String, Object>();
		try {
			ColumnFamilyHandle columnFamilyHandle = groupMap.get(group);
			if (null != columnFamilyHandle) {
				RocksIterator iter = rocksDb.newIterator(columnFamilyHandle);
				for(iter.seekToFirst(); iter.isValid(); iter.next()) {
					map.put(new String(iter.key()), ObjectByte.toObject(iter.value()));
				}
			}
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
		return map;
	}
	
	@Override
	public int count(String group) {
		if (null == group) {
			group = GROUP_DEFAULT;
		}
		int count = 0;
		try {
			ColumnFamilyHandle columnFamilyHandle = groupMap.get(group);
			if (null != columnFamilyHandle) {
				RocksIterator iter = rocksDb.newIterator(columnFamilyHandle);
				for(iter.seekToFirst(); iter.isValid(); iter.next()) {
					count++;
				}
			}
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
		return count;
	}
	
	@Override
	public Object get(String group, String key) {
		if (null == group) {
			group = GROUP_DEFAULT;
		}
		try {
			ColumnFamilyHandle columnFamilyHandle = groupMap.get(group);
			if (null != columnFamilyHandle) {
				return ObjectByte.toObject(
					rocksDb.get(columnFamilyHandle, key.getBytes())
				);
			}
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
		return null;
	}
	
	@Override
	public void remove(String group, String key) {
		if (null == group) {
			group = GROUP_DEFAULT;
		}
		try {
			ColumnFamilyHandle columnFamilyHandle = groupMap.get(group);
			if (null != columnFamilyHandle) {
				rocksDb.delete(columnFamilyHandle, key.getBytes());
//				rocksDb.singleDelete(groupHandleMap.get(group), key.getBytes());
			}
		}
		catch (Throwable e) {
			e.printStackTrace();
		}
	}
	
	@Override
	public String type() {
		return KvFactoryConst.KV_ROCKSDB;
	}

	private static KvRocksDb kvObject = null;
	public static final KvRocksDb getInstance() {
		if (null == kvObject) {
			kvObject = new KvRocksDb();
		}
		return kvObject;
	}

}
